import pandas as pd
import statsmodels.api as sm
import numpy as np
import matplotlib.pyplot as plt
from _code.Utils import *
from _code.Modules import *
import CountRegression

def extraModel(x):
    # return SAEModel.ae2(SAEModel.ae1(x, True), True)
    return mySAE.ae1(x, True)

def extraModel2(x):
    return mySAE.ae2(mySAE.ae1(x, True), True)


varSize = 4
hiddenSize1 = 3
hiddenSize2 = 2
historyRecordNum = 1000

isAE = True
isFilter = False
isTwoLayer = True

mySAE = torch.load("./modelStorage/"+str(varSize)+"to"+str(hiddenSize1)+"to"+str(hiddenSize2)+".pth")


xMat, yMat = getNoisData()

xPredi = torch.tensor(xMat, dtype=torch.float32)
yPredi = torch.tensor(yMat)

# 模型提取特征
if isAE:
    if isTwoLayer:
        xPredi = torch.cat((extraModel(xPredi), extraModel2(xPredi)), 1)
    else:
        xPredi = extraModel(xPredi)


tempRInfo = []
preInfo = xPredi.detach().numpy()
ypreInfo = yPredi.detach().numpy()
for i in range(preInfo.shape[1]):
    temp1 = np.corrcoef(preInfo[:, i], ypreInfo.flat)[0, 1]
    tempRInfo.append(np.abs(temp1) > 0.1)
print(tempRInfo)
# 是否过滤掉低相关性的特征
if isFilter:
    xPredi = xPredi[:, tempRInfo]


X_train = xPredi.detach().numpy()[0:1000, :]
y_train = yPredi.detach().numpy()[0:1000,]
X_test = xPredi.detach().numpy()[1000:, :]
y_test = yPredi.detach().numpy()[1000:,]

print(X_train.shape, y_train.shape)

pModel = sm.Poisson(y_train, X_train+0.01).fit()
preds = pModel.predict(X_test+0.01)
plt.plot(y_test[200:300], c = "blue")
plt.plot(preds[200:300], c = "red")
plt.show()

# pModel = CountRegression.PoissonRegression(X_train,y_train)
# pModel.fit()
# print(pModel.predict(X_test)[:5])
#
# pModel = sm.GLM(y_train, X_train, family=sm.families.Poisson())
# results = pModel.fit()
# print(results.summary())



